package com.kancy.delay.queue.redis.service.impl;

import com.kancy.delay.queue.redis.config.DelayQueueRedisTemplate;
import com.kancy.delay.queue.redis.config.DelayQueueTaskExecutor;
import com.kancy.delay.queue.redis.config.RedisScriptConfig;
import com.kancy.delay.queue.redis.config.TraceContextResolver;
import com.kancy.delay.queue.redis.handler.DelayQueueExceptionHandler;
import com.kancy.delay.queue.redis.handler.DelayQueueMessageHandler;
import com.kancy.delay.queue.redis.handler.RetryDelayQueueExceptionHandler;
import com.kancy.delay.queue.redis.key.RedisKeyResolver;
import com.kancy.delay.queue.redis.message.DelayMessage;
import com.kancy.delay.queue.redis.properties.DelayQueueProperties;
import com.kancy.delay.queue.redis.service.DelayQueueServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * <p>
 * DelayQueueServerImpl
 * <p>
 *
 * @author: kancy
 * @date: 2020/3/4 21:10
 **/
public class DelayQueueServerImpl implements DelayQueueServer, ApplicationContextAware {
    private static final Logger log = LoggerFactory.getLogger(DelayQueueServerImpl.class);

    /**
     * 服务启动标志
     */
    private final AtomicBoolean started = new AtomicBoolean(false);

    /**
     * 处理中
     */
    private final AtomicBoolean processing = new AtomicBoolean(false);

    private final DelayQueueRedisTemplate delayQueueRedisTemplate;
    private final RedisKeyResolver redisKeyResolver;
    private final DelayQueueProperties delayQueueProperties;
    private final DelayQueueExceptionHandler defaultExceptionHandler;
    private final DelayQueueTaskExecutor delayQueueTaskExecutor;

    private ScheduledExecutorService scheduleExecutor;
    private ApplicationContext applicationContext;

    @Autowired(required = false)
    private TraceContextResolver traceContextResolver;

    public DelayQueueServerImpl(DelayQueueRedisTemplate delayQueueRedisTemplate, RedisKeyResolver redisKeyResolver,
                                DelayQueueProperties delayQueueProperties,
                                DelayQueueExceptionHandler defaultExceptionHandler,
                                DelayQueueTaskExecutor delayQueueTaskExecutor) {
        this.delayQueueRedisTemplate = delayQueueRedisTemplate;
        this.redisKeyResolver = redisKeyResolver;
        this.delayQueueProperties = delayQueueProperties;
        this.defaultExceptionHandler = defaultExceptionHandler;
        this.delayQueueTaskExecutor = delayQueueTaskExecutor;
    }

    public DelayQueueServerImpl(DelayQueueRedisTemplate delayQueueRedisTemplate, RedisKeyResolver redisKeyResolver,
                                DelayQueueProperties delayQueueProperties, DelayQueueTaskExecutor delayQueueTaskExecutor) {
        this(delayQueueRedisTemplate,redisKeyResolver,delayQueueProperties,
                new RetryDelayQueueExceptionHandler(delayQueueTaskExecutor, delayQueueProperties), delayQueueTaskExecutor);
    }

    /**
     * 启动Server
     */
    @Override
    public final void startServer() {
        if (delayQueueProperties.getServer().isEnabled() && started.compareAndSet(false, true)){
            if (Objects.isNull(scheduleExecutor)){
                scheduleExecutor = new ScheduledThreadPoolExecutor(1, new CustomizableThreadFactory("delay-scanner-"));
            }
            // 开启调度任务
            scheduleExecutor.scheduleAtFixedRate(this::doServer,
                    delayQueueProperties.getServer().getInitialDelay().getSeconds(),
                    delayQueueProperties.getServer().getScanInterval().getSeconds(), TimeUnit.SECONDS);
            Runtime.getRuntime().addShutdownHook(new Thread(this::shutdown));
            log.info("Delay queue server started.");
        }
    }

    /**
     * 停止Server
     */
    @Override
    public final void stopServer() {
        shutdown();
    }

    /**
     * doServer
     */
    private void doServer() {
        if (!processing.getAndSet(true) && started.get()){
            delayQueueTaskExecutor.execute(() -> {
                // Topic 的处理尽量同时处理，防止前面的Topic处理过慢而导致，后面的Topic被拖延处理
                List<Future<?>> futureList = new ArrayList<>(delayQueueProperties.getTopics().size());
                delayQueueProperties.getTopics().entrySet()
                        .stream()
                        .filter(e -> e.getValue().isEnabled())
                        .forEach(e -> futureList.add(delayQueueTaskExecutor.submit(() -> handleTopic(e.getKey()))));

                // 异步阻塞等到任务处理完成
                futureList.forEach(future -> {
                    try {
                        future.get();
                    } catch (Exception e) {
                        log.error(e.getMessage(), e);
                    }
                });

                // 处理完成，重置，防止任务被重复取出，浪费性能
                // 等到上一次任务处理完毕，才继续从队列中取出任务
                // 线程池线程数不够时，会阻塞当前线程
                processing.set(false);
            });
        }else {
            log.debug("The last task has not been completed, this task is ignored.");
        }
    }

    /**
     * 处理延时队列Topic
     * @param topic
     */
    private void handleTopic(String topic) {
        try {
            Set<Object> messageIds = prefetchMessageId(topic);
            log.debug("开始处理：Delay Queue Topic -> {} , taskNum -> {}", topic, messageIds.size());

            if (!CollectionUtils.isEmpty(messageIds)){
                // 循环处理每一条消息
                // 当可用线程不够时，阻塞主线程
                messageIds.stream().forEach(messageId -> delayQueueTaskExecutor.execute(() -> handleMessage(topic, messageId)));
            }
        } catch (Exception e) {
            log.error(String.format("处理Delay Queue Topic [%s] 失败：%s", topic,e.getMessage()), e);
            try {
                getExceptionHandler(topic).onHandleTopicException(topic, e);
            } catch (Exception ex) {
                log.error("DelayQueueExceptionHandler onHandleTopicException throws exception" , ex);
            }
        }
    }

    /**
     * 预取符合条件的消息id
     * @param topic
     * @return
     */
    private Set<Object> prefetchMessageId(String topic) {
        String bucketKey = redisKeyResolver.bucketKey(topic);
        long maxScore = System.currentTimeMillis();
        int prefetchSize = delayQueueProperties.getTopics().get(topic).getPrefetchSize();

        // ZRANGEBYSCORE key min max [WITHSCORES] [LIMIT offset count]
        // http://redisdoc.com/sorted_set/zrangebyscore.html#
        if (prefetchSize > 0){
            return delayQueueRedisTemplate.opsForZSet().rangeByScore(bucketKey, 0, maxScore, 0, prefetchSize);
        }
        return delayQueueRedisTemplate.opsForZSet().rangeByScore(bucketKey, 0, maxScore);
    }

    /**
     * 处理消息
     * @param topic
     * @param messageId
     */
    private void handleMessage(String topic, Object messageId) {
        String currentTraceId = MDC.get(DelayMessage.TRACE_ID_KEY);
        DelayQueueMessageHandler handler = null;
        DelayMessage delayMessage = null;
        try {
            handler = getDelayQueueMessageHandler(topic);
            delayMessage = getAndRemoveDelayMessage(topic, messageId);
            // 保存traceId，便于在分布式环境中请求链路完整性
            if (Objects.nonNull(traceContextResolver)){
                traceContextResolver.setTraceId(delayMessage.getTraceId());
            }
        } catch (Exception e) {
            log.error(String.format("处理Delay Queue Topic [%s] Message [%s] 失败：%s", topic, messageId, e.getMessage()), e);
            try {
                getExceptionHandler(topic).onHandleMessageException(topic, messageId, e);
            } catch (Exception ex) {
                log.error("DelayQueueExceptionHandler onHandleMessageException throws exception" , ex);
            }
            return;
        }

        try {
            // 注意：因为分任务是异步的，非常快；
            // 如果消费的比较慢就会导致消息阻塞在线程池阻塞队列中
            // 一旦发生宕机，就会导致进程中的消息全部丢失
            // 所以，建议不要进行耗时操作，可以把消息丢到MQ或者Redis队列中再进行消费

            // 针对以上问题，做了些优化，放慢了分任务的速度（2020/03/13）
            // 1. 上一个任务没有处理完成，直接忽略当前任务
            // 2. 任务处理，固定线程池数量，去除阻塞队列缓冲，线程数不足，由主线程执行，从而阻塞主线程，放慢分任务的速度
            if (Objects.nonNull(delayMessage)){
                handler.process(delayMessage);
            }
        } catch (Exception e) {
            // 消息处理异常
            log.error("DelayQueueMessageHandler process topic [{}] delayMessage [{}] fail: {}", topic, messageId, e.getMessage());
            // 交给异常处理器进行重试
            try {
                getExceptionHandler(topic).onProcessMessageException(handler, delayMessage, e);
            } catch (Exception ex) {
                log.error("DelayQueueExceptionHandler onProcessMessageException throws exception" , ex);
            }
        } finally {
            if (Objects.nonNull(traceContextResolver)){
                MDC.put(DelayMessage.TRACE_ID_KEY, currentTraceId);
            }
        }
    }

    /**
     * 获取消息处理器
     * @param topic
     * @return
     */
    private DelayQueueMessageHandler getDelayQueueMessageHandler(String topic) {
        String messageHandlerBeanName = null;
        try {
            messageHandlerBeanName = delayQueueProperties.getTopics().get(topic).getMessageHandler();
            return applicationContext.getBean(messageHandlerBeanName, DelayQueueMessageHandler.class);
        } catch (Exception e) {
            messageHandlerBeanName = String.format("%s%s", messageHandlerBeanName, DelayQueueMessageHandler.class.getSimpleName());
            return applicationContext.getBean(messageHandlerBeanName, DelayQueueMessageHandler.class);
        }
    }

    /**
     * 获取异常处理器
     * @param topic
     * @return
     */
    private DelayQueueExceptionHandler getExceptionHandler(String topic) {
        String exceptionHandlerBeanName = null;
        try {
            exceptionHandlerBeanName = delayQueueProperties.getTopics().get(topic).getExceptionHandler();
            return applicationContext.getBean(exceptionHandlerBeanName, DelayQueueExceptionHandler.class);
        } catch (Exception e) {
            try {
                exceptionHandlerBeanName = String.format("%s%s", exceptionHandlerBeanName, DelayQueueExceptionHandler.class.getSimpleName());
                return applicationContext.getBean(exceptionHandlerBeanName, DelayQueueExceptionHandler.class);
            } catch (Exception ex) {
                return defaultExceptionHandler;
            }
        }
    }


    /**
     * 关闭
     */
    private void shutdown(){
        try {
            if (!scheduleExecutor.isShutdown()){
                // 发出停止信号，使定时任务停止工作
                started.set(false);
                // 关闭线程池
                scheduleExecutor.shutdown();
                scheduleExecutor.awaitTermination(
                        delayQueueProperties.getExecutor().getAwaitTerminationSeconds(), TimeUnit.SECONDS);
                log.info("Delay queue server scheduledThreadPoolExecutor shutdown success.");
            }
        } catch (InterruptedException e) {
            log.warn("Delay queue server scheduledThreadPoolExecutor shutdown be Interrupted.");
            Thread.currentThread().interrupt();
        }
    }

    /**
     * 获取并且删除消息（原子操作）
     * @param topic
     * @param messageId
     * @return
     */
    private DelayMessage getAndRemoveDelayMessage(String topic, Object messageId) {
        List<String> keys = new ArrayList<>();
        keys.add(redisKeyResolver.hashTableKey(topic));
        keys.add(redisKeyResolver.bucketKey(topic));
        keys.add(redisKeyResolver.key(String.valueOf(messageId)));

        List messageList = (List) delayQueueRedisTemplate.execute(
                RedisScriptConfig.getDelayMessageGetAndRemoveRedisScript(), keys, messageId);
        if (!CollectionUtils.isEmpty(messageList)){
            Object message = messageList.get(0);
            if (message instanceof DelayMessage){
                DelayMessage delayMessage = DelayMessage.class.cast(message);
                return delayMessage;
            }else {
                log.error("{} not instanceof {}", message, DelayMessage.class.getName());
            }
        }
        return null;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) {
        this.applicationContext = applicationContext;
    }
}
